# -*- coding:utf-8 -*-

# @Time    : 2024/4/9 20:57
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : test_tts.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

import asyncio
import json
import wave

# pip install websockets
import websockets

start_msg = {"header": {"event_name": "StartTTS"}}
stop_msg = {"header": {"event_name": "StopTTS"}}


async def main():
    uri = "ws://116.196.96.191/speech_editing_huize/streaming"  #公网
    async with websockets.connect(uri) as ws:
        start_msg_str = json.dumps(start_msg, ensure_ascii=False)
        await ws.send(start_msg_str)
        msg = await ws.recv()
        msg_dict = json.loads(msg)
        # {'header': {'status': 200, 'status_message': 'ok', 'event_name': 'TTSStarted'}, 'payload': {'audio_sample_rate': 8000, 'audio_encoding': 'PCM_S', 'bits_per_sample': 16, 'num_channels': 1}}
        print(msg_dict)
        if msg_dict["header"]["status"] != 200:
            print("error")
            return
        payload = msg_dict["payload"]
        sample_rate = payload["audio_sample_rate"]
        bits_per_sample = payload["bits_per_sample"]
        sample_width = bits_per_sample // 8

        task1 = asyncio.create_task(send_msg(ws))
        task2 = asyncio.create_task(recv_msg(ws))
        await task1
        pcm = await task2
        wav_file = "test.wav"
        with wave.open(wav_file, "wb") as wav_obj:
            # (nchannels, sampwidth, framerate, nframes, comptype, compname)
            wav_obj.setparams(
                (1, sample_width, sample_rate, 0, "NONE", "NONE")
            )
            wav_obj.writeframes(pcm)


async def send_msg(ws):
    text_lst = [
        "哎，你好。嗯，你好，我这是慧择保险网的哈，这不新的一年了嘛。咱们呢是为了回馈老客户，呃，给您免费赠送了一个一对一保险咨询服务。那您看我这儿呢，就耽误您个几分钟时间，简单的给您说一下，可以吧？",
    ]
    for text in text_lst:
        msg = {"header": {"event_name": "Text"}, "payload": {"text": text}}
        msg_str = json.dumps(msg, ensure_ascii=False)
        await ws.send(msg_str)
    stop_msg_str = json.dumps(stop_msg, ensure_ascii=False)
    await ws.send(stop_msg_str)

import pyaudio

audio_player = pyaudio.PyAudio()
stream = audio_player.open(format=pyaudio.paInt16,
                           channels=1,
                           rate=8000,
                           output=True)



async def recv_msg(ws):
    pcm = b""
    while True:
        msg = await ws.recv()
        if isinstance(msg, bytes):
            pcm += msg
            data = stream.write(msg)
            print(msg)
        else:
            msg_dict = json.loads(msg)
            print(msg_dict)
            break
    return pcm


if __name__ == "__main__":
    asyncio.run(main())
